#ifndef TRAINER_H
#define TRAINER_H

#include "common_includes.h"
#include "training_level1.h"
#include "training_level2.h"

template<typename T>
class Trainer{
//Todo
//prepare dataset
//thread
private:
	vector<Music<T>> training_data;
	vector<reference_wrapper<Music<T>>> training_data_references;
	unsigned training_data_total_frames;

	T power;
	T best_state_bias;
	TrainingLevel2<Number<T, NUM_TRANSITIONS> , T> trainer;
	Transitions<Number<T, NUM_TRANSITIONS>> current_transitions;
public:
	Trainer()
	: trainer(40){

		best_state_bias = 1;
	}

	~Trainer(){
	}

	Trainer(const Trainer& trainer) = delete;

	int progress(void *instance, const lbfgsfloatval_t *x,
			const lbfgsfloatval_t *g, const lbfgsfloatval_t fx,
			const lbfgsfloatval_t xnorm, const lbfgsfloatval_t gnorm,
			const lbfgsfloatval_t step, int n, int k, int ls){
//		printf("Iteration %d:\n", k);
//		printf("xnorm = %f, gnorm = %f, step = %f\n", xnorm, gnorm, step);
//		printf("\n");
//		cout << "x norm is !" << xnorm << "\n";
		if (k > 20){
			cout << "TERMINATING !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n";
			return -1;}
		return 0;
	}

	lbfgsfloatval_t evaluate(void* instance, const lbfgsfloatval_t* x,
			lbfgsfloatval_t* g, const int n, const lbfgsfloatval_t step){
		lbfgsfloatval_t fx = 0.0;
		Transitions<Number<T, NUM_TRANSITIONS>> transitions;
		for (unsigned i = 0; i < NUM_TRANSITIONS; ++i){
			transitions.compact[i] = exp(Number<T, NUM_TRANSITIONS>(i, x[i]));
		}
		transitions.normalize_compact();
		for(unsigned i = 0; i < NUM_TRANSITIONS; ++i){
			transitions.compact[i] = pow(transitions.compact[i], power);
		}

		transitions.compact_to_full();
		current_transitions = transitions;

		WorkResult<Number<T,NUM_TRANSITIONS>> work_result;
		work_result = trainer.evaluate(training_data_references, transitions, best_state_bias);
//		work_result = trainer.evaluate({training_data[0]}, transitions, best_state_bias);
//		Number<T, NUM_TRANSITIONS> evaluation = trainers[0]->train( training_data[0], transitions);

		fx = -work_result.score_ground_truth.var;
		for (unsigned i = 0; i < NUM_TRANSITIONS; ++i){
			g[i] = -work_result.score_ground_truth[i];
		}

		unsigned total_frame_size = training_data_total_frames;
		cout << "GT : " << work_result.score_ground_truth.var / total_frame_size << "\t ";
		cout << "BP : " << work_result.score_best_path / total_frame_size << "\t ";
		cout << "BS : " << work_result.score_best_states / total_frame_size << "\t ";
		cout << "BP-GT : " << (double)work_result.similarity_bp_gt / total_frame_size << "\t ";
		cout << "BS-GT : " << (double)work_result.similarity_bs_gt / total_frame_size << "\t ";
		cout << "BP-BS : " << (double)work_result.similarity_bp_bs / total_frame_size << "\t ";
		cout << endl;
		return fx;
	}

	void prepare(string training_data_path, T power = 1, T best_state_bias=1){
		training_data.clear();
		training_data_references.clear();
		training_data_total_frames = 0;
		namespace fs = boost::filesystem;
		fs::path someDir(training_data_path);
		fs::directory_iterator end_iter;

		unsigned count = 500;
		training_data.reserve(500); //Todo, how about shared_pointer instead of reference_wrapper
		for( fs::directory_iterator dir_iter(someDir); dir_iter != end_iter; ++dir_iter){

			if (fs::is_regular_file(dir_iter->status()) && count<-1){
				cout << dir_iter->path() << endl;
				training_data.push_back(Music<T>(dir_iter->path().string().c_str()));
				training_data_references.push_back(training_data.back());
				training_data_total_frames += training_data.back().frame_size;
			}
			if (count++ == 16){
				break;
			}
		}

		this->best_state_bias = best_state_bias;
		this->power = power;
		for (auto i = training_data.begin(); i != training_data.end(); ++i){
			for (auto j = i->observation.begin(); j != i->observation.end(); ++j){
				for (auto k = j->begin(); k != j->end(); ++k){
					*k = pow(*k, power);
				}
			}
		}
	}

	void run(){
		int ret = -2;
		lbfgsfloatval_t fx;
		lbfgsfloatval_t *x = lbfgs_malloc(NUM_TRANSITIONS);
		lbfgs_parameter_t param;
		lbfgs_parameter_init(&param);
		for (unsigned i = 0; i < NUM_TRANSITIONS; ++i){
			x[i] = 1;
		}
		while (true){
			switch (ret){
				case -1:
				for (unsigned i = 0; i < NUM_TRANSITIONS; ++i){
					x[i] = log(current_transitions.compact[i].var);
				}
			}
			ret = lbfgs(NUM_TRANSITIONS, x, &fx, evaluate_wrapper, progress_wrapper,
					this, &param);
		}
		cout << "termination " << ret;
	}

	static int progress_wrapper(void *instance, const lbfgsfloatval_t *x,
			const lbfgsfloatval_t *g, const lbfgsfloatval_t fx,
			const lbfgsfloatval_t xnorm, const lbfgsfloatval_t gnorm,
			const lbfgsfloatval_t step, int n, int k, int ls){
		return static_cast<Trainer*>(instance)->progress(instance, x, g, fx, xnorm,
				gnorm, step, n, k, ls);
	}

	static lbfgsfloatval_t evaluate_wrapper(void* instance,
			const lbfgsfloatval_t* x, lbfgsfloatval_t* g, const int n,
			const lbfgsfloatval_t step){
		return static_cast<Trainer*>(instance)->evaluate(instance, x, g, n, step);
	}
};

#endif

